package leetcode_1101_1200;

import node.TreeNode;

import java.util.*;

public class LeeCode_1110 {
    public static void main(String[] args) {
        System.out.println(delNodes(new TreeNode(1, new TreeNode(2), new TreeNode(3, null, new TreeNode(4))), new int[]{1, 2}));
    }

    private static List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        Set<Integer> set = new HashSet<>();
        for (int num : to_delete) {
            set.add(num);
        }
        List<TreeNode> ans = new ArrayList<>();
        ans.add(root);
        TreeNode dummy = new TreeNode(0);
        dfs(ans, root, dummy, set);
        List<TreeNode> nodes = new ArrayList<>();
        for (TreeNode node : ans) {
            if (set.contains(node.val)) {
                nodes.add(node);
            }
        }
        for (TreeNode node : nodes) {
            ans.remove(node);
        }
        return ans;
    }

    private static void dfs(List<TreeNode> ans, TreeNode node, TreeNode pre, Set<Integer> set) {
        if (node == null)
            return;
        if (set.contains(node.val)) {
            if (pre.left != null && pre.left == node){
                pre.left = null;
            }else {
                pre.right = null;
            }
            if (node.left != null) {
                ans.add(node.left);
            }
            if (node.right != null) {
                ans.add(node.right);
            }
        }
        dfs(ans, node.left, node, set);
        dfs(ans, node.right, node, set);
    }
}
